#!/usr/bin/env python
################################################################################
#
#   svg_from_transcripts
# 
#
#
#   Copyright (C) 2007 Leo Goodstadt
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; version 2
#   of the License
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################
"""
    svg_from_transcripts.py

"""

import sys, os
import os.path
import re
import StringIO
import glob
import shutil
# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(exe_path)

# get module name
if __name__ == '__main__':
    myname = os.path.splitext(os.path.split(sys.argv[0])[1])[0];
else:
    myname = __name__

from SVGdraw import *
import SVGdraw
from general_util import *


def parse_exon_data(file):
    transcripts = dict()
    cnt_lines = 0
    for line in file:
        cnt_lines += 1
        line = line.rstrip()
        if not len(line):
            continue;
        if line[0] == '#':
            continue
        try:
            transcript, start, end =  line.split("\t", 3)
        except ValueError:
            die_error("Need three fields on line #%d\n[%s]" % (cnt_lines, line))
        start = int(start)
        end = int(end)
        transcripts.setdefault(transcript, []).append([start, end])
    return transcripts

from plot_scaling import plot_scaling 

def plot_transcripts_from_file (scaling, transcript_exon_data, exon_dimensions, 
                                    wobbly_pos, straight_pos, text_pos, draw_properties):
    transcript_data = parse_exon_data(transcript_exon_data)
    return plot_transcripts (scaling, transcript_data, exon_dimensions, wobbly_pos, 
                    straight_pos, text_pos, draw_properties)

from svg_draw_properties import svg_draw_properties
    
def plot_transcripts (scaling, transcripts, exon_dimensions, wobbly_pos, 
                    straight_pos, text_pos, draw_properties):
    """
        draws "wobbly line stringing together exons.
            with wobbly_pos = y coordinate of the extent of the wobbly line
            with straight_pos = y coordinate of the exonic (straight part)
          if wobbly_pos < straight_pos, the line will be drawn pointing down
            
        exon_dimensions is a list of the op and bottom of the exon rectangles

        scaling is a type of plot_scaling

        transcripts is a dictionary of lists of exonic data 
        first two fields of exonic data must be beg and ends
    """

    # hold end result
    svg_elements = dict()
    for element in ['exons', 'line', 'name']:
        svg_elements[element] = dict()
    
    for element in ['exons', 'line', 'name']:
        if element not in draw_properties:
            draw_properties[element] = svg_draw_properties()
            
    wobbly_line = list()
    wobbly_pos = scaling.to_y(wobbly_pos)
    straight_pos = scaling.to_y(straight_pos)

    # exons
    exon_top = scaling.to_y(exon_dimensions[0])
    exon_bot = scaling.to_y(exon_dimensions[1])

    prev_exon_end = 0
    for transcript_name in transcripts:
        exon_group = group(transcript_name)
        exons = list(transcripts[transcript_name])
        exons.sort()
        for exonic_data in exons:
            beg, end = exonic_data [0:2]

            # draw genomic line
            if prev_exon_end:
                # draw intron
                midpoint = (beg + prev_exon_end) / 2
                wobbly_line.append([scaling.to_x(midpoint), wobbly_pos])

            # add point for exon
            wobbly_line.append([scaling.to_x(beg), straight_pos])
            wobbly_line.append([scaling.to_x(end), straight_pos])

            #
            #   draw exons
            # 
            c = rect(scaling.to_x(beg),
                     exon_bot, 
                     scaling.x_scaled(end - beg + 1), 
                     exon_top - exon_bot, 
                     **draw_properties['exons']())
            exon_group.addElement(c)
            prev_exon_end = end
        prev_exon_end = 0

        svg_elements['exons'][transcript_name] = exon_group

        # 
        # draw wobbly line
        # 
        if len(wobbly_line) > 2:
            #pl=polyline(wobbly_line, stroke = 'blue', stroke_width = wobby_width, fill = 
            #            'none')
            #exon_group.addElement(pl)
            pd=pathdata(wobbly_line[0][0], wobbly_line[0][1])
            for x, y in wobbly_line[1:]:
                pd.line(x, y)
            svg_elements['line'][transcript_name] = path(pd, **draw_properties['line']())

        #
        #   draw text
        # 
        svg_elements['name'][transcript_name] = text(x = (wobbly_line[0][0] + 
                                                          wobbly_line[-1][0]) / 2.0, 
                                                    y = scaling.to_y(text_pos), 
                                                    text = transcript_name,
                                                    text_anchor = 'middle',
                                                     **draw_properties['name']())
        wobbly_line = list()

            
    return svg_elements

if __name__ == '__main__':
    test_data = \
"""transcript_0	64488451	64489651
transcript_1	65392735	65392892
transcript_1	65384020	65384148
transcript_2	65579997	65580192
transcript_2	65588259	65588338
transcript_2	65589265	65589549
transcript_3	65668727	65668823
transcript_3	65653312	65653457
transcript_3	65669018	65669115
transcript_3	65679098	65679357
transcript_4	65764193	65764316
transcript_5	65771073	65771148
transcript_6	65824229	65824340
transcript_7	66062478	66062733
transcript_8	66101527	66101761
transcript_8	66169092	66169220
transcript_8	66110652	66110792
transcript_8	66262473	66262616
transcript_8	66261277	66262222
transcript_8	66151000	66151115
transcript_8	66120072	66120232
transcript_8	66257208	66257322
transcript_8	66171788	66171982
"""
    test_data = test_data.split("\n")
    # 
    # create drawing then you create a SVG root element
    # 
    drawing=drawing()
    svg_root=svg()

    sys.stderr.write("plot exons...\n")
    
    # in mm
    paper_width  = 297
    paper_height = 210
    paper_margin = 20
    genomic_extent = 2000000.0
    genomic_offset = 66473839.0
    internal_width  = paper_width - paper_margin * 2
    internal_height = paper_height - paper_margin * 2


    # svg root discription
    svg_root.attributes["width"]    = str(paper_height) + "mm"
    svg_root.attributes["height"]   = str(paper_width) + "mm"
    svg_root.attributes["viewBox"] = "0 0 %d %d" % (paper_width, paper_height)
    transform_group = group("root transform")
    transform_group.attributes["transform"] = "translate(%d,%d)" % (0, paper_height)
    svg_root.addElement(transform_group)

    #
    # scaling factor
    # 
    genome_scaling = plot_scaling(-internal_width/genomic_extent, -20, 
                                  paper_margin, paper_margin, 
                                    genomic_offset, 0 )


    #
    # colour of lines
    # 
    draw_properties = dict()
    draw_properties['line'] = svg_draw_properties(stroke = 'red', 
                                                        stroke_width = 0.2)
    draw_properties['exons'] = svg_draw_properties(stroke = 'None', 
                                                    fill = 'green')
    draw_properties['name'] = svg_draw_properties(stroke = 'None', 
                                                    fill = 'blue',
                                                    font_family="Arial",
                                                    font_size = abs(genome_scaling.y_scaled(0.5)))

    svg_elements = plot_transcripts_from_file (genome_scaling, test_data, 
                                                    [-0.5, 0.5], -0.6, -0.5, 2.0,
                                                    draw_properties)


    exon_group = group("all_exons")
    lines_group = group("all_lines")
    names_group = group("all_text")
    for e in svg_elements['exons'].values():
        exon_group.addElement(e)
    for e in svg_elements['line'].values():
        lines_group.addElement(e)
    for e in svg_elements['name'].values():
        names_group.addElement(e)
        
    #
    # add exons
    # 
    transform_group.addElement(exon_group)
    transform_group.addElement(lines_group)
    transform_group.addElement(names_group)
    
    drawing.setSVG(svg_root)
    print drawing.toXml()
    sys.stderr.write("Done...\n")

